package com.sanjiayi.websocket;

import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONObject;
import com.sanjiayi.pojo.ChatMessage;
import com.sanjiayi.service.ChatService;
import com.sanjiayi.utils.JwtUtil;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;

import javax.websocket.*;
import javax.websocket.server.PathParam;
import javax.websocket.server.ServerEndpoint;
import java.io.IOException;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

@Slf4j
@ServerEndpoint("/ws/chat/{userId}/{token}")
@Component
public class ChatWebSocketServer {
    
    private static ChatService chatService;
    private static JwtUtil jwtUtil;
    // 存储在线用户的WebSocket会话
    private static final Map<Integer, Session> onlineUsers = new ConcurrentHashMap<>();
    // 存储正在通话中的用户
    private static final Map<Integer, Integer> inCallUsers = new ConcurrentHashMap<>();
    
    @Autowired
    public void setChatService(ChatService chatService) {
        ChatWebSocketServer.chatService = chatService;
    }
    
    @Autowired
    public void setJwtUtil(JwtUtil jwtUtil) {
        ChatWebSocketServer.jwtUtil = jwtUtil;
    }
    
    @OnOpen
    public void onOpen(Session session, @PathParam("userId") Integer userId, @PathParam("token") String token) {
        log.info("User {} trying to connect", userId);
        
        if (!jwtUtil.validateToken(token)) {
            try {
                session.close(new CloseReason(CloseReason.CloseCodes.VIOLATED_POLICY, "Invalid token"));
            } catch (IOException e) {
                log.error("Close session error", e);
            }
            return;
        }

        // 如果用户已经在线，先断开旧连接
        Session oldSession = onlineUsers.get(userId);
        if (oldSession != null && oldSession.isOpen()) {
            try {
                oldSession.close();
            } catch (IOException e) {
                log.error("Close old session error", e);
            }
        }
        
        onlineUsers.put(userId, session);
        log.info("User {} connected, current online users: {}", userId, onlineUsers.size());
    }
    
    @OnClose
    public void onClose(@PathParam("userId") Integer userId) {
        onlineUsers.remove(userId);
        // 如果用户正在通话中，通知对方通话结束
        Integer callPartnerId = inCallUsers.remove(userId);
        if (callPartnerId != null) {
            sendCallEndedNotification(callPartnerId, userId);
        }
        log.info("User {} disconnected, current online users: {}", userId, onlineUsers.size());
    }
    
    @OnMessage
    public void onMessage(String message, @PathParam("userId") Integer userId) {
        log.info("Received message from user {}: {}", userId, message);
        try {
            JSONObject jsonMessage = JSON.parseObject(message);
            Integer receiverId = jsonMessage.getInteger("receiverId");
            String type = jsonMessage.getString("type");
            
            if (receiverId == null) {
                log.error("Message missing receiverId: {}", message);
                sendErrorMessage(userId, "消息缺少接收者ID");
                return;
            }

            // 添加发送者ID
            jsonMessage.put("senderId", userId);
            
            // 获取接收者的session
            Session receiverSession = onlineUsers.get(receiverId);
            
            if (receiverSession != null && receiverSession.isOpen()) {
                String messageStr = jsonMessage.toJSONString();
                
                try {
                    // 处理视频通话相关的消息
                    if (handleVideoCallMessage(type, userId, receiverId, jsonMessage)) {
                        receiverSession.getBasicRemote().sendText(messageStr);
                        return;
                    }
                    
                    // 处理普通聊天消息
                    receiverSession.getBasicRemote().sendText(messageStr);
                    if (type == null || "TEXT".equals(type)) {
                        ChatMessage chatMessage = new ChatMessage();
                        chatMessage.setSenderId(userId);
                        chatMessage.setReceiverId(receiverId);
                        chatMessage.setContent(jsonMessage.getString("content"));
                        chatMessage.setMessageType("TEXT");
                        chatService.saveMessage(chatMessage);
                    }
                } catch (IOException e) {
                    log.error("Send message to receiver {} failed", receiverId, e);
                    sendErrorMessage(userId, "消息发送失败");
                    
                    if (!receiverSession.isOpen()) {
                        onlineUsers.remove(receiverId);
                        log.info("Removed inactive session for user {}", receiverId);
                    }
                }
            } else {
                log.warn("Receiver {} is offline or session closed", receiverId);
                sendErrorMessage(userId, "接收者不在线");
                
                if (receiverSession != null && !receiverSession.isOpen()) {
                    onlineUsers.remove(receiverId);
                    log.info("Removed closed session for user {}", receiverId);
                }
            }
        } catch (Exception e) {
            log.error("Process message error", e);
            sendErrorMessage(userId, "消息处理失败");
        }
    }

    /**
     * 处理视频通话相关的消息
     * @param type 消息类型
     * @param senderId 发送者ID
     * @param receiverId 接收者ID
     * @param message 消息内容
     * @return 是否是视频通话消息
     */
    private boolean handleVideoCallMessage(String type, Integer senderId, Integer receiverId, JSONObject message) {
        switch (type) {
            case "call-request":
                // 检查发起者和接收者是否都在线且未在通话中
                if (isUserAvailableForCall(senderId) && isUserAvailableForCall(receiverId)) {
                    log.info("Call request from {} to {}", senderId, receiverId);
                    return true;
                } else {
                    sendErrorMessage(senderId, "对方正在通话中或不在线");
                    return false;
                }
                
            case "call-accept":
                // 建立通话连接
                inCallUsers.put(senderId, receiverId);
                inCallUsers.put(receiverId, senderId);
                log.info("Call accepted between {} and {}", senderId, receiverId);
                return true;
                
            case "call-reject":
                log.info("Call rejected by {} to {}", senderId, receiverId);
                return true;
                
            case "call-end":
                // 结束通话
                inCallUsers.remove(senderId);
                inCallUsers.remove(receiverId);
                log.info("Call ended between {} and {}", senderId, receiverId);
                return true;
                
            case "ice-candidate":
            case "offer":
            case "answer":
                // WebRTC信令消息直接转发
                log.debug("Forwarding WebRTC signal: {} from {} to {}", type, senderId, receiverId);
                return true;
                
            default:
                return false;
        }
    }

    /**
     * 检查用户是否可以进行通话
     * @param userId 用户ID
     * @return 是否可以通话
     */
    private boolean isUserAvailableForCall(Integer userId) {
        return onlineUsers.containsKey(userId) && !inCallUsers.containsKey(userId);
    }

    /**
     * 发送通话结束通知
     * @param userId 用户ID
     * @param partnerUserId 通话对方的用户ID
     */
    private void sendCallEndedNotification(Integer userId, Integer partnerUserId) {
        try {
            Session session = onlineUsers.get(userId);
            if (session != null && session.isOpen()) {
                JSONObject notification = new JSONObject();
                notification.put("type", "call-end");
                notification.put("senderId", partnerUserId);
                notification.put("reason", "对方已断开连接");
                session.getBasicRemote().sendText(notification.toJSONString());
            }
        } catch (IOException e) {
            log.error("Send call ended notification failed", e);
        }
    }
    
    @OnError
    public void onError(Session session, Throwable error) {
        log.error("WebSocket error: {}", error.getMessage());
    }
    
    private void sendErrorMessage(Integer userId, String errorMessage) {
        try {
            Session session = onlineUsers.get(userId);
            if (session != null && session.isOpen()) {
                JSONObject error = new JSONObject();
                error.put("type", "error");
                error.put("message", errorMessage);
                session.getBasicRemote().sendText(error.toJSONString());
            }
        } catch (IOException e) {
            log.error("Send error message failed", e);
        }
    }
    
    public static void sendMessage(Integer userId, String message) {
        Session session = onlineUsers.get(userId);
        if (session != null && session.isOpen()) {
            try {
                session.getBasicRemote().sendText(message);
            } catch (IOException e) {
                log.error("Send message error", e);
            }
        }
    }
} 